from scipy.optimize import minimize
from scipy.optimize import LinearConstraint
import numpy as np

def func_adc_raw(arg,y,z,w):
    return np.sum((arg[0]+arg[1]*y+arg[2]*z-w)**2)

def func_sigma_m_raw(sigma,m,w,x):
    y = (x-m)/sigma
    z = np.sqrt(y**2+1)
    
    func_adc = lambda arg: func_adc_raw(arg,y,z,w)
    lc = LinearConstraint([[1,0,0],[0,0,1],[0,1,1],[0,-1,1]],[0,0,0,0],[max(w),4*sigma,4*sigma,4*sigma])
    opt_out = minimize(func_adc,x0=[0,0,0],method='SLSQP',constraints=lc)

    return opt_out.fun,opt_out.x

def calibration_absrm(w,x,max_sigma = 10,method='Nelder-Mead'):
    # sigma : 0 - 10
    # m     : min(x) - max(x)
    region = [[0.01,max_sigma],[min(x),max(x)]]
    func_sigma_m = lambda arg: func_sigma_m_raw(arg[0],arg[1],w,x)[0]
    
    if method == 'Nelder-Mead':
        opt_out = minimize(func_sigma_m,x0=[0.01,min(x)],method='Nelder-Mead',\
            bounds=region,options={'fatol':1e-10,'xatol':1e-5})
        sigma,m = opt_out.x
    elif method == 'TR':
        opt_out = minimize(func_sigma_m,x0=[0.01,min(x)],method='trust-constr',\
            bounds=region)
        sigma,m = opt_out.x
    else:
        assert 1<0,"Calib adjust.py: Invalid Method"

    a , d, c = func_sigma_m_raw(sigma,m,w,x)[1]
    return a, c/sigma, sigma, d/c, m # a b sigma rho m